#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <sm2_sign.h>
#include <sm3.h>
#include <utils.h>

//公私钥生成,私钥签名，公钥验签
int sm2_key_generator(SM2_KEY *key) {
    // private_key的取值是范围在[1, n-2]的随机数
    do {
        if (sm2_bn_rand_range(key->private_key, sm2_bn_order_minus_one()) != 1) {
            print_log();
            return -1;
        }
    } while (sm2_bn_is_zero(key->private_key));
    sm2_bn_point_mul_generator(&key->public_key, key->private_key);
    return 1;
}

//公私钥打印
int sm2_key_printf(const char *label, const SM2_KEY *key) {
    sm2_printf("%s ", label);
    sm2_bn_point_printf("publicKey", &key->public_key);
    sm2_printf("%s ", label);
    sm2_bn_printf("privateKey", key->private_key);
    return 1;
}


// 预计算 d' = (d + 1)^-1 (mod n)
int sm2_sign_compute_key(const SM2_KEY *key, sm2_bn_type fast_private) {
    if (sm2_bn_compare(key->private_key, sm2_bn_order_minus_one()) >= 0) {
        print_log();
        return -1;
    }
    sm2_bn_mod_n_plus(fast_private, key->private_key, sm2_bn_one());
    sm2_bn_mod_n_inv(fast_private, fast_private);
    return 1;
}

void sm2_bn_point_fast_get_xy(SM2_SIGN_PRE_COMP *pre_comp, sm2_bn_point_type *P) {
    sm2_bn_type f[SM2_SIGN_PRE_COMP_COUNT];
    sm2_bn_type g[SM2_SIGN_PRE_COMP_COUNT];
    int i;
    //下面以SM2_SIGN_PRE_COMP_COUNT=32为例
    // f[0] = Z[0]
    // f[1] = Z[0] * Z[1]
    // ...
    // f[31] = Z[0] * Z[1] * ... * Z[31]
    sm2_bn_copy(f[0], P[0].Z);
    for (i = 1; i < SM2_SIGN_PRE_COMP_COUNT; i++) {
        sm2_bn_mod_p_mont_multiplication(f[i], f[i - 1], P[i].Z);
    }

    // 在这里只用进行一次求模逆的运算f[31]^-1 = (Z[0] * ... * Z[31])^-1
    sm2_bn_mod_p_mont_inv(f[SM2_SIGN_PRE_COMP_COUNT - 1], f[SM2_SIGN_PRE_COMP_COUNT - 1]);

    // g[31] = Z[31]
    // g[30] = Z[30] * Z[31]
    // ...
    // g[1] = Z[1] * Z[2] * ... * Z[31]
    //
    sm2_bn_copy(g[SM2_SIGN_PRE_COMP_COUNT - 1], P[SM2_SIGN_PRE_COMP_COUNT - 1].Z);
    for (i = SM2_SIGN_PRE_COMP_COUNT - 2; i >= 1; i--) {
        sm2_bn_mod_p_mont_multiplication(g[i], g[i + 1], P[i].Z);
    }
    // 累乘快速得到Z[i]^-1
    // Z[0]^-1 = g[1] * f[31]^-1
    // Z[1]^-1 = g[2] * f[0] * f[31]^-1
    // ...
    // Z[31]^-1 = f[30] * f[31]^-1
    sm2_bn_mod_p_mont_multiplication(P[0].Z, g[1], f[SM2_SIGN_PRE_COMP_COUNT - 1]);
    for (i = 1; i <= SM2_SIGN_PRE_COMP_COUNT - 2; i++) {
        sm2_bn_mod_p_mont_multiplication(P[i].Z, g[i + 1], f[i - 1]);
        sm2_bn_mod_p_mont_multiplication(P[i].Z, P[i].Z, f[SM2_SIGN_PRE_COMP_COUNT - 1]);
    }
    sm2_bn_mod_p_mont_multiplication(P[SM2_SIGN_PRE_COMP_COUNT - 1].Z, f[SM2_SIGN_PRE_COMP_COUNT - 2],
                                     f[SM2_SIGN_PRE_COMP_COUNT - 1]);

    // 快速得到需要的x[i],x[i] = X[i] * Z[i]^-2 (mod n)
    for (i = 0; i < SM2_SIGN_PRE_COMP_COUNT; i++) {
        sm2_bn_mod_p_mont_sqr(P[i].Z, P[i].Z);
        sm2_bn_mod_p_mont_multiplication(pre_comp[i].x1_mod_n, P[i].X, P[i].Z);
        sm2_bn_mod_p_from_mont(pre_comp[i].x1_mod_n, pre_comp[i].x1_mod_n);
        if (sm2_bn_compare(pre_comp[i].x1_mod_n, sm2_bn_order()) >= 0) {
            sm2_bn_subtraction(pre_comp[i].x1_mod_n, pre_comp[i].x1_mod_n, sm2_bn_order());
        }
    }
}

// 使用Montgomery's 的技巧在多个(x1， y1) = [k]G 上得到Z的逆坐标
int sm2_sign_pre_compute(SM2_SIGN_PRE_COMP pre_comp[SM2_SIGN_PRE_COMP_COUNT]) {
    sm2_bn_point_type P[SM2_SIGN_PRE_COMP_COUNT];
    for (int i = 0; i < SM2_SIGN_PRE_COMP_COUNT; i++) {
        // rand k in [1, n - 1]
        do {
            if (sm2_bn_rand_range(pre_comp[i].k, sm2_bn_order()) != 1) {
                print_log();
                return -1;
            }
        } while (sm2_bn_is_zero(pre_comp[i].k));

        // (x1, y1) = kG窗口法解决
        sm2_bn_point_mul_generator(&P[i], pre_comp[i].k);
    }
    //小技巧快速得到x1
    sm2_bn_point_fast_get_xy(pre_comp, P);
    return 1;
}


// s = (k - r * d)/(1 + d)
//	= (k + r)*(1 + d)^-1 - r
//	= (k + r) * d' - r, d'是预计算得到的结果
int sm2_do_sign(const sm2_bn_type fast_private, const SM2_SIGN_PRE_COMP *pre_comp, const uint8_t digest[32],
                SM2_SIGNATURE *sig) {
    sm2_bn_type e, r, s;
    // 得到摘要 e = H(M)
    sm2_bn_from_bytes(e, digest);
    if (sm2_bn_compare(e, sm2_bn_order()) >= 0) {
        sm2_bn_subtraction(e, e, sm2_bn_order());
    }

    // r = e + x1 (mod n)
    sm2_bn_mod_n_plus(r, e, pre_comp->x1_mod_n);

    // s = (k + r) * d' - r
    sm2_bn_mod_n_plus(s, pre_comp->k, r);
    sm2_bn_mod_n_to_mont(s, s);
    sm2_bn_mod_n_mont_multiplication(s, s, fast_private); // mont(s) * d = s * R^-1 * d * R = s * d
    sm2_bn_mod_n_subtraction(s, s, r);

    sm2_bn_to_bytes(r, sig->r);
    sm2_bn_to_bytes(s, sig->s);
    return 1;
}

int sm2_do_verify(const sm2_bn_point_type point_table[16], const uint8_t digest[32], const SM2_SIGNATURE *sig) {
    sm2_bn_point_type R, T;
    sm2_bn_type r, s, e, x, t;

    // 检查 r, s 是否属于 [1, n-1]
    sm2_bn_from_bytes(r, sig->r);
    if (sm2_bn_is_zero(r) == 1) {
        print_log();
        return -1;
    }
    if (sm2_bn_compare(r, sm2_bn_order()) >= 0) {
        print_log();
        return -1;
    }
    sm2_bn_from_bytes(s, sig->s);
    if (sm2_bn_is_zero(s) == 1) {
        print_log();
        return -1;
    }
    if (sm2_bn_compare(s, sm2_bn_order()) >= 0) {
        print_log();
        return -1;
    }

    // t = r + s (mod n), check t != 0
    sm2_bn_mod_n_plus(t, r, s);
    if (sm2_bn_is_zero(t)) {
        print_log();
        return -1;
    }

    // Q(x,y) = s * G + t * P
    sm2_bn_point_mul_generator(&R, s);
    sm2_bn_point_mul_ex(&T, t, point_table);
    sm2_bn_point_plus(&R, &R, &T);
    sm2_bn_point_get_xy(&R, x, NULL);

    // e = H(M)
    sm2_bn_from_bytes(e, digest);
    if (sm2_bn_compare(e, sm2_bn_order()) >= 0) {
        sm2_bn_subtraction(e, e, sm2_bn_order());
    }

    // r' = e + x (mod n)
    if (sm2_bn_compare(x, sm2_bn_order()) >= 0) {
        sm2_bn_subtraction(x, x, sm2_bn_order());
    }
    sm2_bn_mod_n_plus(e, e, x);

    // 验证是否 r == r'
    if (sm2_bn_compare(e, r) != 0) {
        print_log();
        return -1;
    }
    return 1;
}

int sm2_signature_to_der(const SM2_SIGNATURE *sig, uint8_t **out, size_t *out_len) {
    size_t len = 0;
    if (!sig) {
        return 0;
    }
    if (asn1_integer_2_der_ex(sig->r, 32, NULL, &len) != 1
        || asn1_integer_2_der_ex(sig->s, 32, NULL, &len) != 1
        || asn1_header_to_der(len, out, out_len) != 1
        || asn1_integer_2_der_ex(sig->r, 32, out, out_len) != 1
        || asn1_integer_2_der_ex(sig->s, 32, out, out_len) != 1) {
        print_log();
        return -1;
    }
    return 1;
}

int sm2_printf_signature(const char *label, const SM2_SIGNATURE *sig) {
    sm2_printf("%s ", label);
    sm2_printf_bytes("r", sig->r, 32);
    sm2_printf("%s ", label);
    sm2_printf_bytes("s", sig->s, 32);
    return 1;
}

int sm2_compute_z(uint8_t z[32], const sm2_bn_point_type *pub, const char *id, const size_t id_len) {
    SM3_CTX ctx;
    uint8_t zin[18 + 32 * 6] = {
        0x00, 0x80,
        0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 0x38, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 0x38,
        0xFF, 0xFF, 0xFF, 0xFE, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
        0xFF, 0xFF, 0xFF, 0xFF, 0x00, 0x00, 0x00, 0x00, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFC,
        0x28, 0xE9, 0xFA, 0x9E, 0x9D, 0x9F, 0x5E, 0x34, 0x4D, 0x5A, 0x9E, 0x4B, 0xCF, 0x65, 0x09, 0xA7,
        0xF3, 0x97, 0x89, 0xF5, 0x15, 0xAB, 0x8F, 0x92, 0xDD, 0xBC, 0xBD, 0x41, 0x4D, 0x94, 0x0E, 0x93,
        0x32, 0xC4, 0xAE, 0x2C, 0x1F, 0x19, 0x81, 0x19, 0x5F, 0x99, 0x04, 0x46, 0x6A, 0x39, 0xC9, 0x94,
        0x8F, 0xE3, 0x0B, 0xBF, 0xF2, 0x66, 0x0B, 0xE1, 0x71, 0x5A, 0x45, 0x89, 0x33, 0x4C, 0x74, 0xC7,
        0xBC, 0x37, 0x36, 0xA2, 0xF4, 0xF6, 0x77, 0x9C, 0x59, 0xBD, 0xCE, 0xE3, 0x6B, 0x69, 0x21, 0x53,
        0xD0, 0xA9, 0x87, 0x7C, 0xC6, 0x2A, 0x47, 0x40, 0x02, 0xDF, 0x32, 0xE5, 0x21, 0x39, 0xF0, 0xA0,
    };

    if (!z || !pub || !id) {
        print_log();
        return -1;
    }

    sm2_bn_point_to_bytes(pub, &zin[18 + 32 * 4]);

    sm3_init(&ctx);
    if (strcmp(id, SM2_DEFAULT_ID) == 0) {
        sm3_update(&ctx, zin, sizeof(zin));
    } else {
        uint8_t id_bits[2];
        id_bits[0] = (uint8_t) (id_len >> 5);
        id_bits[1] = (uint8_t) (id_len << 3);
        sm3_update(&ctx, id_bits, 2);
        sm3_update(&ctx, (uint8_t *) id, id_len);
        sm3_update(&ctx, zin + 18, 32 * 6);
    }
    sm3_finish(&ctx, z);
    return 1;
}

int sm2_signature(SM2_SIGN_CTX *ctx, const SM2_KEY *key, const char *id, const size_t id_len, const uint8_t *data,
                  const size_t data_len, uint8_t *sig, size_t *sig_len, SM2_SIGNATURE *signature) {
    if (!ctx || !key || !sig || !sig_len) {
        print_log();
        return -1;
    }
    // initial
    sm3_init(&ctx->sm3_ctx);
    if (id) {
        uint8_t z[SM3_DIGEST_SIZE];
        if (id_len <= 0 || id_len > SM2_MAX_ID_LENGTH) {
            print_log();
            return -1;
        }
        sm2_compute_z(z, &key->public_key, id, id_len);
        sm3_update(&ctx->sm3_ctx, z, sizeof(z));
    }
    ctx->saved_sm3_ctx = ctx->sm3_ctx;
    if (sm2_sign_pre_compute(ctx->pre_comp) != 1) {//本次实现这里做了优化，但是考虑到现在不需要测速，所以先写了个简介版本
        print_log();
        return -1;
    }
    ctx->num_pre_comp = SM2_SIGN_PRE_COMP_COUNT;
    ctx->key = *key;
    sm2_sign_compute_key(key, ctx->fast_sign_private);

    // update
    if (data && data_len > 0) {
        sm3_update(&ctx->sm3_ctx, data, data_len);
    }

    // do
    uint8_t digest[SM3_DIGEST_SIZE];
    sm3_finish(&ctx->sm3_ctx, digest);
    sm2_printf_bytes("sm2 digest", digest, SM3_DIGEST_SIZE);

    ctx->num_pre_comp--;
    sm2_do_sign(ctx->fast_sign_private, &ctx->pre_comp[ctx->num_pre_comp], digest, signature);
    *sig_len = 0;
    if (sm2_signature_to_der(signature, &sig, sig_len) != 1) {
        print_log();
        return -1;
    }
    return 1;
}

int sm2_verify(SM2_VERIFY_CTX *ctx, const SM2_KEY *key, const char *id, const size_t id_len, const uint8_t *data,
               const size_t data_len, const SM2_SIGNATURE *signature) {
    if (!ctx || !key) {
        print_log();
        return -1;
    }

    sm3_init(&ctx->sm3_ctx);
    if (id) {
        uint8_t z[SM3_DIGEST_SIZE];
        if (id_len <= 0 || id_len > SM2_MAX_ID_LENGTH) {
            print_log();
            return -1;
        }
        sm2_compute_z(z, &key->public_key, id, id_len);
        sm3_update(&ctx->sm3_ctx, z, sizeof(z));
    }
    ctx->saved_sm3_ctx = ctx->sm3_ctx;

    if (!&key->public_key) {
        //取公钥
        print_log();
        return -1;
    }
    ctx->key.public_key = key->public_key;
    sm2_bn_set_zero(ctx->key.private_key);
    sm2_bn_point_mul_pre_compute(&key->public_key, ctx->public_point_table);

    if (data && data_len > 0) {
        sm3_update(&ctx->sm3_ctx, data, data_len);
    }
    // do
    uint8_t digest[SM3_DIGEST_SIZE];
    sm3_finish(&ctx->sm3_ctx, digest);
    if (sm2_do_verify(ctx->public_point_table, digest, signature) != 1) {
        print_log();
        return -1;
    }
    return 1;
}
